from copy import deepcopy
import sys
sys.path.append("../lib")
import unitree_arm_interface
import numpy as np
import matplotlib.pyplot as plt
import copy
import rospy
import time 
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped, Quaternion
from std_msgs.msg import Float64MultiArray
from visualization_msgs.msg import Marker
import tf.transformations as tfs

def pubObstacle(ob_pub, type,mid, x,y,z,sx,sy,sz,tx="1"):
    shape = []
    if type==0:
        shape = Marker.CUBE
    elif type==1:
        shape = Marker.SPHERE
    elif type==2:
        shape = Marker.ARROW
    elif type==3:
        shape = Marker.CYLINDER
    else:
        shape = Marker.TEXT_VIEW_FACING
    marker = Marker()  # 定义Marker类型的变量
    marker.header.frame_id = "world"  # 设置header的frame——id和stamp属性值
    marker.header.stamp = rospy.Time.now()  # 这里应该是一个浮点值，对应cpp版本 ros::Time::now()
    # 下面两个量定义了该marker的命名空间和id， 通过这两个量能够确定marker的身份，如果身份相同，那么则刷新rviz中对应目标。后面我们会讲到marker是什么
    marker.ns = "basic_shapes"
    marker.id = mid
    
    marker.type = shape  # marker的类型
    marker.action = Marker.ADD # marker action

    # 初始化了marker的位置和初始的角度
    marker.pose.position.x = x
    marker.pose.position.y = y
    marker.pose.position.z = z
    marker.pose.orientation.x = 0.0
    marker.pose.orientation.y = 0.0
    marker.pose.orientation.z = 0.0
    marker.pose.orientation.w = 1.0
    # marker的尺寸大小，这里1.0对应于现实地图的1m
    marker.scale.x = sx
    marker.scale.y = sy
    marker.scale.z = sz
    # marker的颜色定义， 注意透明度a的设置，为0就看不到了
    marker.color.r = 155/255.0
    marker.color.g = 155/255.0
    marker.color.b = 155/255.0
    marker.color.a = 1.0

    marker.lifetime = rospy.Duration()
    ob_pub.publish(marker)  # 发布消息
    
    # 每次发布的消息不同，从而让rviz绘制不同的marker


rospy.init_node('APF_ARM_DYNAMIC')
np.set_printoptions(precision=3, suppress=True)
arm =  unitree_arm_interface.ArmInterface(hasGripper = True)
armModel = arm._ctrlComp.armModel

# path = [[ 0,0,0, 0.39       ,  0.29        ,     0.22      ],
#         [ 0,0,0, 0.209, -0.107, 0.308],
#         [ 0,0,0, 0.45       , -0.29        ,     0.33      ]]
# path = np.array(path)

path = np.load("patharray.npy")  # 读取文件
I = np.zeros(path.shape,dtype= np.int)
path = np.concatenate((I, path), axis=1)

q_k = path[0]
path = path[1:path.shape[0]]
temp_goal = path[0]
path = path[1:path.shape[0]]

goal = path[-1]
c_time = 0.033
EE_start = np.asarray([
    [ 1  ,   0. ,    0 ,  0.39],
 [    0    , 1     ,  0.  ,   0.29 ],
 [    0 ,    0  ,    1 , 0.22],
 [    0.     , 0. ,    0  ,  1.   ]])

# print(tfs.euler_from_matrix(EE_start))
      
hasIK, theta_forward_k = armModel.inverseKinematics(EE_start, np.zeros(6), True)
count = 0
cost = [np.linalg.norm(q_k-goal)]

real_path = [q_k]
refrencepath = [q_k]

joint1 = [theta_forward_k[0]]
joint2 = [theta_forward_k[1]]
joint3 = [theta_forward_k[2]]
joint4 = [theta_forward_k[3]]
joint5 = [theta_forward_k[4]]
joint6 = [theta_forward_k[5]]

while np.linalg.norm(q_k-goal)>0.005 and count<1000:
    # print("count ",count)
    J = armModel.CalcJacobian(theta_forward_k)
    # J = J[3:6,:]
    # print("J ",J)
    q_k_d = temp_goal - q_k
    # print("q_k_d1  ",q_k_d)
    q_k_d = q_k_d/np.linalg.norm(q_k_d)/10
    # print(np.linalg.norm(q_k_d))
    # print("q_k_d2  ",q_k_d)

    q_k_d = q_k_d.reshape(6,1)
    # print("q_k_d ",q_k_d)
    theta_forward_dot = np.linalg.pinv(J).dot(q_k_d)
    theta_forward_dot = theta_forward_dot.flatten()
    theta_forward_k = theta_forward_k + theta_forward_dot*c_time
    # print("theta_forward_k2",theta_forward_k)

    joint1.append(theta_forward_k[0])
    joint2.append(theta_forward_k[1])
    joint3.append(theta_forward_k[2])
    joint4.append(theta_forward_k[3])
    joint5.append(theta_forward_k[4])
    joint6.append(theta_forward_k[5])

    q_k = armModel.forwardKinematics(theta_forward_k, 6)
    el = tfs.euler_from_matrix(q_k)
    q_k = np.array([el[0],el[1],el[2], q_k[0,3],q_k[1,3],q_k[2,3]])
    # print("q_k",q_k)

    cost.append(np.linalg.norm(q_k-goal))
    real_path.append(copy.deepcopy(q_k))
    
    if np.linalg.norm(q_k-temp_goal)<0.005 and path.shape[0]!=0:
        temp_goal = path[0]
        path = path[1:path.shape[0]]
    count = count+1
    
rate = rospy.Rate(30)
pub = rospy.Publisher('/joint_group_position_controller/command', Float64MultiArray, queue_size=10)
count = len(joint1)
i = 0


ob_pub = rospy.Publisher("visualization_marker", Marker, queue_size=1)
dya = 0.42    

for j in range(10):
    array = [joint1[i],joint2[i],joint3[i],joint4[i],joint5[i],joint6[i]]
    left_top = Float64MultiArray(data=array)
    pub.publish(left_top)
    time.sleep(0.1)
    
    pubObstacle(ob_pub,1,0, dya,0,0.28,0.1,0.1,0.1)

path_pub = rospy.Publisher('trajectory', Path, queue_size=50)
darw_path = np.load("patharray.npy") 
path_record = Path()
current_time = rospy.Time.now()
path_record.header.frame_id = "world"
path_record.header.stamp = current_time
while not rospy.is_shutdown():
    if i<count:
        array = [joint1[i],joint2[i],joint3[i],joint4[i],joint5[i],joint6[i]]
        left_top = Float64MultiArray(data=array)
        pub.publish(left_top)
        
        pubObstacle(ob_pub,1,0, dya,0,0.28,0.1,0.1,0.1)
        if i < 160:
            dya = dya - 0.05/30.0
        current_time = rospy.Time.now()
        pose = PoseStamped()
        pose.header.stamp = current_time
        pose.header.frame_id = 'world'
        pose.pose.position.x = real_path[i][3]
        pose.pose.position.y = real_path[i][4]
        pose.pose.position.z = real_path[i][5]
        pose.pose.orientation.x = 0
        pose.pose.orientation.y = 0
        pose.pose.orientation.z = 0
        pose.pose.orientation.w = 1
        path_record.poses.append(pose)
        path_pub.publish(path_record)   
        i = i+1        
    rate.sleep()    
        


# from copy import deepcopy
# import sys
# sys.path.append("../lib")
# import unitree_arm_interface
# import numpy as np
# import matplotlib.pyplot as plt
# import copy
# import rospy
# from nav_msgs.msg import Path
# from geometry_msgs.msg import PoseStamped, Quaternion
# from std_msgs.msg import Float64MultiArray
# import tf.transformations as tfs

# rospy.init_node('APF_ARM_DYNAMIC')
# np.set_printoptions(precision=3, suppress=True)
# arm =  unitree_arm_interface.ArmInterface(hasGripper = True)
# armModel = arm._ctrlComp.armModel

# path = [[ 0.39       ,  0.29        ,     0.22      ],
#         [ 0.17197539 , -0.10697421  ,     0.33107302],
#         [ 0.41623353 , -0.27776082  ,     0.33820648],
#         [ 0.45       , -0.29        ,     0.33      ]]
# path = np.array(path)

# q_k = path[0]
# path = path[1:path.shape[0]]
# temp_goal = path[0]
# path = path[1:path.shape[0]]

# goal = path[-1]
# c_time = 0.033
# EE_start = np.asarray([
#     [ 1  ,   0. ,    0 ,  0.39],
#  [    0    , 1     ,  0.  ,   0.29 ],
#  [    0 ,    0  ,    1 , 0.22],
#  [    0.     , 0. ,    0  ,  1.   ]])

# # print(tfs.euler_from_matrix(EE_start))
      
# hasIK, theta_forward_k = armModel.inverseKinematics(EE_start, np.zeros(6), True)
# count = 0
# cost = [np.linalg.norm(q_k-goal)]

# real_path = [q_k]
# refrencepath = [q_k]

# joint1 = [theta_forward_k[0]]
# joint2 = [theta_forward_k[1]]
# joint3 = [theta_forward_k[2]]
# joint4 = [theta_forward_k[3]]
# joint5 = [theta_forward_k[4]]
# joint6 = [theta_forward_k[5]]

# while np.linalg.norm(q_k-goal)>0.005 and count<1000:
#     # print("count ",count)
#     J = armModel.CalcJacobian(theta_forward_k)
#     J = J[3:6,:]
#     # print("J ",J)
#     q_k_d = temp_goal - q_k
#     # print("q_k_d1  ",q_k_d)
#     q_k_d = q_k_d/np.linalg.norm(q_k_d)/10
#     # print(np.linalg.norm(q_k_d))
#     # print("q_k_d2  ",q_k_d)

#     q_k_d = q_k_d.reshape(3,1)
#     # print("q_k_d ",q_k_d)
#     theta_forward_dot = np.linalg.pinv(J).dot(q_k_d)
#     theta_forward_dot = theta_forward_dot.flatten()
#     theta_forward_k = theta_forward_k + theta_forward_dot*c_time
#     # print("theta_forward_k2",theta_forward_k)

#     joint1.append(theta_forward_k[0])
#     joint2.append(theta_forward_k[1])
#     joint3.append(theta_forward_k[2])
#     joint4.append(theta_forward_k[3])
#     joint5.append(theta_forward_k[4])
#     joint6.append(theta_forward_k[5])

#     q_k = armModel.forwardKinematics(theta_forward_k, 6)
#     q_k = np.array([q_k[0,3],q_k[1,3],q_k[2,3]])
#     # print("q_k",q_k)

#     cost.append(np.linalg.norm(q_k-goal))
#     real_path.append(copy.deepcopy(q_k))
    
#     if np.linalg.norm(q_k-temp_goal)<0.005 and path.shape[0]!=0:
#         temp_goal = path[0]
#         path = path[1:path.shape[0]]
#     count = count+1
    
# rate = rospy.Rate(30)
# pub = rospy.Publisher('/joint_group_position_controller/command', Float64MultiArray, queue_size=10)
# count = len(joint1)
# i = 0
# while not rospy.is_shutdown():
#     if i<count:
#         array = [joint1[i],joint2[i],joint3[i],joint4[i],joint5[i],joint6[i]]
#         left_top = Float64MultiArray(data=array)
#         pub.publish(left_top)
#         i = i+1
#         rate.sleep()    